"""
Stain normalization inspired by method of:

A. Vahadane et al., ‘Structure-Preserving Color Normalization and Sparse Stain Separation for Histological Images’, IEEE Transactions on Medical Imaging, vol. 35, no. 8, pp. 1962–1971, Aug. 2016.

"""

from __future__ import division

import os

import cv2
import spams
import numpy as np
from tqdm import tqdm
import shutil

import stain_utils as ut


def get_stain_matrix(I, threshold=0.8, lamda=0.1):
    """
    Get 2x3 stain matrix. First row H and second row E
    :param I:
    :param threshold:
    :param lamda:
    :return:
    """
    mask = ut.notwhite_mask(I, thresh=threshold).reshape((-1,))
    OD = ut.RGB_to_OD(I).reshape((-1, 3))
    OD = OD[mask]
    if OD.size == 0:
        raise ValueError("all pixels have all been masked as being to bright")
    dictionary = spams.trainDL(OD.T, K=2, lambda1=lamda, mode=2, modeD=0, posAlpha=True, posD=True, verbose=False).T
    if dictionary[0, 0] < dictionary[1, 0]:
        dictionary = dictionary[[1, 0], :]
    dictionary = ut.normalize_rows(dictionary)
    return dictionary


class Normalizer(object):
    """
    A stain normalization object
    """

    def __init__(self):
        self.stain_matrix_target = None

    def fit(self, target):
        target = ut.standardize_brightness(target)
        self.stain_matrix_target = get_stain_matrix(target)

    def target_stains(self):
        return ut.OD_to_RGB(self.stain_matrix_target)

    def transform(self, I):
        I = ut.standardize_brightness(I)
        stain_matrix_source = get_stain_matrix(I)
        source_concentrations = ut.get_concentrations(I, stain_matrix_source)
        return (255 * np.exp(-1 * np.dot(source_concentrations, self.stain_matrix_target).reshape(I.shape))).astype(
            np.uint8)

    def hematoxylin(self, I):
        I = ut.standardize_brightness(I)
        h, w, c = I.shape
        stain_matrix_source = get_stain_matrix(I)
        source_concentrations = ut.get_concentrations(I, stain_matrix_source)
        H = source_concentrations[:, 0].reshape(h, w)
        H = np.exp(-1 * H)
        return H


if __name__ == '__main__':
    image_dir = "/media/hsmy/wanghao_18T/dataset/MIDOG2021/tiff/"

    output_dir = "/media/hsmy/wanghao_18T/dataset/5fold/fold1/tiff_aug/"
    os.makedirs(output_dir, exist_ok=True)
    domain_dir = "/media/hsmy/wanghao_18T/dataset/5fold/domain/"

    for domain_file in os.listdir(domain_dir):
        domain_name = domain_file.split(".")[0]
        im_path = os.path.join(domain_dir, domain_file)
        im = cv2.imread(im_path)
        im = cv2.cvtColor(im, cv2.COLOR_BGR2RGB)

        normalizer = Normalizer()
        normalizer.fit(im)
        for input_file in tqdm(os.listdir(image_dir)):
            img_name = input_file.split(".")[0]

            filter_arr = [118, 123, 39, 143, 11, 109, 117, 29, 125, 149, 76, 116, 46, 144, 38, 5, 82, 84, 128, 16, 139, 120, 124, 88, 23, 114, 70,14, 57, 112, 3, 122, 93, 132, 74, 85, 22, 12, 89, 113, 66, 2, 140, 40, 32, 45, 60, 87, 54, 36, 145, 105, 63, 96,102, 92, 75, 49, 95, 97, 73, 134, 43, 25, 100, 9, 41, 119, 127, 78, 62, 71, 64, 107, 31, 130, 91, 58, 17, 136, 21,80, 90, 99, 142, 47, 103, 69, 81, 24, 30, 131, 72, 79, 104, 137, 138, 55, 108, 13, 8, 110, 27, 34, 106, 111, 65, 56]
            if int(img_name) not in filter_arr:
                continue

            file_path = os.path.join(image_dir, input_file)
            img = cv2.imread(file_path)
            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            transformed_img = normalizer.transform(img)

            file_name = img_name + "_" + domain_name

            cv2.imwrite(output_dir + file_name + ".tiff",
                        cv2.cvtColor(transformed_img, cv2.COLOR_RGB2BGR))

        print(f"domain {domain_name} done")
